import os
import platform
import sys
import sysconfig

from setuptools import setup, Extension, find_packages, find_namespace_packages
from setuptools.command.build_ext import build_ext as build_ext_orig

python_paths = sysconfig.get_paths()
print(f'python_paths: {python_paths}')

library_dirs = []
py_limited_api = True

if sys.platform == 'win32':
    v = sys.version_info
    if v.minor == 10:
        version = '3.10.11'
    elif v.minor == 11:
        version = '3.11.9'
    elif v.minor == 12:
        version = '3.12.9'
    elif v.minor == 13:
        version = '3.13.2'
    else:
        raise ValueError('不知道的 Python 版本')

    arch = os.getenv('CIBW_ARCHS_WINDOWS', 'AMD64')
    print(f'arch: {arch}')
    if arch == 'AMD64':
        arch = 'x64'
    python_lib = os.path.join(r'C:\hostedtoolcache\windows\Python', version, arch, 'libs')
    print(f'python_lib: {python_lib} {os.path.exists(python_lib)} {os.walk(python_lib)}')
    os.environ['CC'] = 'gcc'
    os.environ['LD'] = 'ld'
    extra_compile_args = ['-O3', '-fopenmp', '-static']
    extra_link_args = ['-static', '-fopenmp', '-Wl,-Bstatic', '-lgomp']
    library_dirs = [python_lib]
    py_limited_api = False
elif sys.platform == 'darwin':
    os.environ['CC'] = 'gcc-14'
    os.environ['MACOSX_DEPLOYMENT_TARGET'] = '12.0'
    libomp_prefix = os.environ.get('LIBOMP_PREFIX', '/opt/homebrew/opt/libomp')
    extra_compile_args = [
        '-fopenmp',
        f'-I{libomp_prefix}/include'
    ]
    extra_link_args = [
        f'-L{libomp_prefix}/lib',
        '-fopenmp',
        '-lomp',
        f'-Wl,-rpath,{libomp_prefix}/lib'
    ]
    if platform.machine() == 'arm64':
        # 针对 ARM 架构（Apple Silicon）
        extra_compile_args.append('-march=armv8-a')
else:
    extra_compile_args = ['-fopenmp']
    extra_link_args = ['-fopenmp']


class CustomBuildExt(build_ext_orig):
    def build_extensions(self):
        # 检查系统类型，自动配置编译器
        if os.name == 'nt':
            from distutils import ccompiler
            compiler = ccompiler.new_compiler(compiler='mingw32')
            os.environ['CC'] = 'gcc'
            os.environ['CXX'] = 'g++'
            self.compiler = compiler
        super().build_extensions()


ext_modules = [
    Extension(
        'cqlib.simulator.statevector_simulator_c',
        include_dirs=[
            'cqlib/simulator/c_fusion/include',
            python_paths['include'],
        ],
        library_dirs=[python_paths['platlib']] + library_dirs,  # 链接库文件所在目录
        sources=[
            'cqlib/simulator/c_fusion/src/c_fusion.c',
            'cqlib/simulator/c_fusion/src/state_vector.c',
        ],
        define_macros=[('Py_LIMITED_API', '0x030A0000')] if py_limited_api else [],
        py_limited_api=py_limited_api,
        extra_compile_args=extra_compile_args + (['-DPy_LIMITED_API=0x030A0000'] if py_limited_api else []),
        extra_link_args=extra_link_args,
    ),
]

print(f"find_namespace_packages: {find_namespace_packages(include=['cqlib', 'cqlib.*'])}")
print(f"find_packages: {find_packages(include=['cqlib', 'cqlib.*'])}")
setup(
    name='cqlib',
    ext_modules=ext_modules,
    packages=find_namespace_packages(include=['cqlib', 'cqlib.*']),
    package_data={
        'cqlib': ['*'],
    },
    include_package_data=True,
    python_requires='>=3.10, <4',
    exclude_package_data={
        'cqlib.simulator': ['*.c', '*.h'],
    },
    cmdclass={'build_ext': CustomBuildExt},
    options={"bdist_wheel": {"py_limited_api": "cp310"} if py_limited_api else {}},
)
